# -*- coding: utf-8 -*-
# created on 20176/8/23
import copy

from mathsolver.functions.base import *
from mathsolver.functions.shulie.Fundamental import *


class ShuLieANameSNameUpdate001(BaseFunction):
    """
    更新数列名
    input: a_(n), S_(n)
    eg. 设公比为q的等比数列\\{a_{n}\\}的前n项和为S_{n},若S_{2}=3a_{2}+2,S_{4}=3a_{4}+2,q>0,则a_{1}=( )
    """
    def solver(self, *args):
        assert len(args) == 2
        an = args[0].sympify()
        sn = args[1].sympify()
        an_names = sl_shulie_name(an)
        assert len(an_names) == 1
        an_name = an_names[0]
        sn_names = sl_shulie_name(sn)
        assert len(sn_names) == 1
        sn_name = sn_names[0]
        sl = self.search(str(an_name))
        if sl:
            sl.sName = str(sn_name)
            name = str(an_name) + str(sn_name)
            sl.name = name
            del self.known[str(an_name)]
            self.output.append(sl)
        if sl.type == "DengCha":
            self.label.add("等差数列的前n项和")
        elif sl.type == "DengBi":
            self.label.add("等比数列的前n项和")
        else:
            self.label.add("数列的前n项和")
        return self


class ShuLieANameSNameUpdate002(BaseFunction):
    """
    更新数列名
    input: S_(n)
    eg. 设正项等比数列{a_{n}}的首项a_{1}=\\frac{1}{2},前n项和为S_{n},2^{10}S_{30}-(2^{10}+1)S_{20}+S_{10}=0,求{a_{n}}的通项.
    """
    def solver(self, *args):
        assert len(args) == 1
        sn = args[0].sympify()
        sls = [self.known[item] for item in self.known if isinstance(self.known[item], BaseSequence)]
        assert len(sls) == 1
        sl = sls[0]
        ex_name = sl.name
        s_names = sl_shulie_name(sn)
        s_names = [str(t) for t in s_names]
        s_names = [t for t in s_names if t != sl.aName]
        s_name = s_names[0]
        sl.sName = s_name
        a_name = sl.aName
        name = a_name + s_name
        sl.name = name
        del self.known[ex_name]
        self.output.append(sl)
        if sl.type == "DengCha":
            self.label.add("等差数列的前n项和")
        elif sl.type == "DengBi":
            self.label.add("等比数列的前n项和")
        else:
            self.label.add("数列的前n项和")
        return self


class ShuLieANameSNameUpdate003(BaseFunction):
    """
    更新数列名
    input: a_(n), eq/eqs
    eg. 若数列{a_{n}}的前n项和S_{n}=\\dfrac{2}{3}a_{n}+\\dfrac{1}{3},则{a_{n}}的通项公式是a_{n}=.
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSequence)
        assert isinstance(args[1], (BaseSeqEq, BaseEq, BaseSeqEqs, BaseEqs))
        sl = self.known[args[0].name]
        if isinstance(args[1], (BaseSeqEq, BaseEq)):
            eq = args[1].sympify()
            expr = eq[0] - eq[1]
            p = r'[A-Z]_'
            m = re.findall(p, str(expr))
            m = list(set(m))
            if len(m) == 1:
                s_name = m[0]
                sl.sName = s_name
                a_name = sl.aName
                name = a_name + s_name
                sl.name = name
                del self.known[args[0].name]
                self.output.append(sl)
        if sl.type == "DengCha":
            self.label.add("等差数列的前n项和")
        elif sl.type == "DengBi":
            self.label.add("等比数列的前n项和")
        else:
            self.label.add("数列的前n项和")
        return self


class ShuLieANameSNameUpdate004(BaseFunction):
    """
    更新数列名
    input: eq/eqs
    eg. 已知数列{a_{n}}中,a_{1}=1,前n项和{S}_{n}=\\frac{n+2}{3}{a}_{n},求a_{10}
    """
    def solver(self, *args):
        sls = [self.known[item] for item in self.known if isinstance(self.known[item], BaseSequence)]
        assert len(sls) == 1
        assert len(args) == 1
        assert isinstance(args[0], (BaseSeqEq, BaseEq, BaseSeqEqs, BaseEqs))
        sl = sls[0]
        if isinstance(args[0], (BaseSeqEq, BaseEq)):
            eq = args[0].sympify()
            expr = eq[0] - eq[1]
            p = r'[A-Z]_'
            m = re.findall(p, str(expr))
            m = list(set(m))
            if len(m) == 1:
                del self.known[sl.name]
                s_name = m[0]
                sl.sName = s_name
                a_name = sl.aName
                name = a_name + s_name
                sl.name = name
                self.output.append(sl)
        if sl.type == "DengCha":
            self.label.add("等差数列的前n项和")
        elif sl.type == "DengBi":
            self.label.add("等比数列的前n项和")
        else:
            self.label.add("数列的前n项和")
        return self


class ShuLieANameSNameUpdate(BaseFunction):
    CLS = [ShuLieANameSNameUpdate001, ShuLieANameSNameUpdate002,
           ShuLieANameSNameUpdate003, ShuLieANameSNameUpdate004]

    def solver(self, *args):
        r = None
        for cl in ShuLieANameSNameUpdate.CLS:
            try:
                r = cl(self.known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r
